# Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved.
import sys
from datetime import datetime

import pytz
from sqlalchemy import BigInteger, Column, DateTime, Integer, String, create_engine, BOOLEAN
from sqlalchemy.orm import declarative_base, sessionmaker

from golden_seed_server.content.config import config

Base = declarative_base()


class User(Base):
    __tablename__ = "user"
    __table_args__ = {
        "mysql_engine": "InnoDB",
        "mysql_charset": "utf8mb4",
        "mysql_collate": "utf8mb4_general_ci"
    }
    id = Column(BigInteger().with_variant(Integer, "sqlite"), primary_key=True, autoincrement=True)
    user_sub = Column(String(length=100))
    passwd = Column(String(length=100), nullable=True)
    organization = Column(String(length=100), nullable=True)
    revision_number = Column(String(length=100), nullable=True)
    login_time = Column(DateTime, nullable=True)
    created_time = Column(DateTime, default=lambda: datetime.now(pytz.timezone('Asia/Shanghai')))
    credit = Column(Integer, default=100)
    is_whitelisted = Column(BOOLEAN, default=False)


class MysqlDB:
    _engine = None

    @classmethod
    def get_mysql_engine(cls):
        if not cls._engine:
            cls.engine = create_engine(
                    f'mysql+pymysql://{config["MYSQL_USER"]}:{config["MYSQL_PWD"]}'
                    f'@{config["MYSQL_HOST"]}:{config["MYSQL_PORT"]}/{config["MYSQL_DATABASE"]}',
                    hide_parameters=True,
                    echo=False,
                    pool_recycle=300,
                    pool_pre_ping=True)
            Base.metadata.create_all(cls.engine)
        return cls._engine

    @classmethod
    def get_session(cls):
        connection = None
        connection = sessionmaker(bind=cls.engine)()
        return cls._ConnectionManager(connection)

    class _ConnectionManager:
        def __init__(self, connection):
            self.connection = connection

        def __enter__(self):
            return self.connection

        def __exit__(self, exc_type, exc_val, exc_tb):
            try:
                self.connection.close()
            except Exception as e:
                MysqlDB.logger.error(f"Mysql connection close failed due to error: {e}")


MysqlDB.get_mysql_engine()
try:
    with MysqlDB.get_session() as session:
        session.query(1)
        pass
except:
    sys.exit(1)


